# ==============================================================================
# 文件: data_loader.py
# 作用: 负责加载原始数据集（如CIFAR-10），并根据配置将其转换为
#      二分类数据，最终模拟成PU（Positive-Unlabeled）数据集。
# ==============================================================================
import torch
from torchvision import datasets, transforms # 从torchvision导入数据集和图像变换工具
from torch.utils.data import Subset, TensorDataset # 导入数据子集和张量数据集工具
import numpy as np # 导入numpy库
import config # 从同级目录的config模块导入配置

def get_and_prepare_cifar10():
    """
    下载CIFAR-10数据集，并根据config中的设定筛选出指定的正负类别。
    :return: 一个包含筛选后图像和对应二元标签(0/1)的元组，(positive_data, negative_data)。
    """
    print("正在加载和准备CIFAR-10数据集...")
    # 定义图像的预处理流程，这里只转换为Tensor
    transform = transforms.Compose([transforms.ToTensor()])

    # 下载训练集和测试集
    train_set = datasets.CIFAR10(root=config.DATA_DIR, train=True, download=True, transform=transform)
    test_set = datasets.CIFAR10(root=config.DATA_DIR, train=False, download=True, transform=transform)

    # 合并训练集和测试集以获得更多数据
    full_dataset = torch.utils.data.ConcatDataset([train_set, test_set])

    # 筛选出指定的正例和负例
    positive_data = [] # 创建一个空列表来存放筛选后的数据
    negative_data = []
    for image, label in full_dataset:
        if label == config.POSITIVE_CLASS:
            # 如果是指定的正类别，标签重映射为1
            positive_data.append((image, 1))
        elif label == config.NEGATIVE_CLASS:
            # 如果是指定的负类别，标签重映射为0
            negative_data.append((image, 0))
    
    print(f"数据集准备完成。正例(类别{config.POSITIVE_CLASS})数量: {len(positive_data)}")
    print(f"负例(类别{config.NEGATIVE_CLASS})数量: {len(negative_data)}")
    return positive_data, negative_data

def create_pu_dataset(positive_data, negative_data):
    """
    将标准的二分类数据集模拟成PU数据集。
    :param positive_data: 包含(图像, 标签)元组的正例列表。
    :param negative_data: 包含(图像, 标签)元组的负例列表。
    :return: P_dataset (仅含正例), U_dataset (未标记混合数据), U_labels_for_eval (U集的真实标签，仅供评估使用)
    """
    print("正在创建PU数据集...")
    # 将数据按标签分离
    positive_samples = [d for d, l in positive_data]
    negative_samples = [d for d, l in negative_data]
    
    # 对正样本进行随机排序，以确保每次抽样都是随机的
    np.random.shuffle(positive_samples)

    # 创建P集 (Positive set)
    P_set_images = positive_samples[:config.N_POSITIVE]
    
    # 创建U集 (Unlabeled set)
    # U集包含剩余的正样本和所有的负样本
    U_set_images = positive_samples[config.N_POSITIVE:] + negative_samples
    # 为U集创建对应的真实标签，这部分信息在训练中不可见，仅用于最终评估
    U_labels_for_eval = [1] * len(positive_samples[config.N_POSITIVE:]) + [0] * len(negative_samples)

    # 为了方便后续处理，将图像列表转换为TensorDataset
    # 注意：这里的P_dataset和U_dataset只包含图像，没有标签
    P_dataset = TensorDataset(torch.stack(P_set_images))
    U_dataset = TensorDataset(torch.stack(U_set_images))
    
    print(f"PU数据集创建完成。")
    print(f"P集大小 (已知正例): {len(P_dataset)}")
    print(f"U集大小 (未标记数据): {len(U_dataset)}")
    print(f"U集中真实正例比例 (pi): {sum(U_labels_for_eval) / len(U_labels_for_eval):.4f}")

    return P_dataset, U_dataset, torch.tensor(U_labels_for_eval, dtype=torch.float32)
